import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

# ============= MNIST数据集探索 ==============
# 读取MNIST数据集
mnist = tf.keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

# 特征数据归一化
train_images = train_images/255.0
test_images = test_images/255.0
# 对标签数据进行独热编码
train_labels_ohe = tf.one_hot(train_labels, depth=10)
test_labels_ohe = tf.one_hot(test_labels, depth=10)

# ============= 模型定义 ===============
# 建立线性堆叠模型
model = tf.keras.models.Sequential()
# 添加平坦层
# 原始输入的数据，一个样本不管shape如何，都将其展平为一维
# 如本案例中，原始输入的shape为(60000,28,28)，展平为(60000,)
model.add(tf.keras.layers.Flatten(input_shape=(28, 28)))
# 添加全连接层1
model.add(tf.keras.layers.Dense(units=64,  # 神经元个数
                                kernel_initializer='normal',  # 参数初始化器
                                activation='relu'))  # 激活函数
# 添加全连接层2
model.add(tf.keras.layers.Dense(units=32,
                                kernel_initializer='normal',
                                activation='relu'))
# 添加输出层
model.add(tf.keras.layers.Dense(units=10, activation='softmax'))
# 输出模型摘要
model.summary()
# 注：以上模型也可以一次性完成
# model = tf.keras.models.Sequential([
#     tf.keras.layers.Flatten(input_shape=(28, 28)),
#     tf.keras.layers.Dense(64, activation=tf.nn.relu),
#     tf.keras.layers.Dense(32, activation=tf.nn.relu),
#     tf.keras.layers.Dense(10, activation=tf.nn.softmax)
# ])
# ================= 定义训练模式 ===================
model.compile(optimizer='adam',  # 优化器
              loss=tf.keras.losses.categorical_crossentropy,  # 损失函数
              metrics=['accuracy'])  # 评估指标

# 设置训练参数
training_epochs = 10  # 训练轮数
batch_size = 30  # 单次训练样本数

# 模型训练
train_history = model.fit(x=train_images, y=train_labels_ohe,
          validation_split=0.2,
          epochs=training_epochs,
          batch_size=batch_size,
          verbose=2)

# 训练过程指标数据
print(train_history.history)
# 训练过程指标可视化
def show_train_history(train_history, train_metric, val_metric):
    plt.plot(train_history.history[train_metric])
    plt.plot(train_history.history[val_metric])
    plt.title('Train History')
    plt.ylabel(train_metric)
    plt.xlabel('Epoch')
    plt.legend(['train', 'validation'], loc='upper left')
    plt.show()

show_train_history(train_history, 'loss', 'val_loss')
show_train_history(train_history, 'accuracy', 'val_accuracy')

# ================= 评估模型 ===================
test_loss, test_acc = model.evaluate(test_images, test_labels_ohe, verbose=2)
print('评估模型')
print(test_loss, test_acc)

# ================= 应用模型 ===================
print('应用模型')
test_pred = model.predict(test_images)
print(test_pred.shape)
print('预测值：', np.argmax(test_pred[0]))
print('实际值：', test_labels[0])
